#!/usr/bin/python3
# -*- coding: utf-8 -*-
import os
import cv2
import struct
import multiprocessing
import time
# import jpeg4py as jpeg

base_path = r'/rootfs/media/kasim/DataSet/FaceVideo'
out_path = r'/rootfs/media/kasim/DataSet/FaceVideo/CropImageMargin'
# bbox_count_file_path = r'/rootfs/media/kasim/DataSet/FaceVideo/video_list_bbox_count.txt'
# bbox_file_path = r'/rootfs/media/kasim/DataSet/FaceVideo/video_list_bbox.dat'
bbox_count_file_path = r'/rootfs/media/kasim/DataSet/FaceVideo/face_video_list_bbox_count.txt'
bbox_file_path = r'/rootfs/media/kasim/DataSet/FaceVideo/face_video_list_bbox.dat'

PROCESS_COUNT = 12

SLEEP_TIME = 0
SHOW = False
SAVE_IMAGE = True
# SHOW = True
# SAVE_IMAGE = False


select_class = 'Face'

select_class_label = {
    'Person': 0,
    'Cat': 1,
    'Dog': 2,
    'BabyCar': 3,
    'Face': 4,
}
select_label = select_class_label[select_class]

THRESHOLDS = [
    0.3,  # Person
    0.3,  # Cat
    0.3,  # Dog
    0.3,  # BabyCar
    0.3,  # Face
]


max_size = 1920
# margins = [1.0, 1.0, 1.0, 1.0]
margins = [0.3, 0.4, 0.3, 0.2]


def get_box(img, result, score_thr=0.3, show=False, color=(0, 255, 0), class_name=''):
    bboxes = []
    for bbox in result:
        if bbox[4] < score_thr:
            continue
        bboxes.append(bbox)
    if len(bboxes) < 1:
        return bboxes
    height = img.shape[0]
    width = img.shape[1]
    ret_bboxes = []
    for i, bbox in enumerate(bboxes):
        left = int(round(bbox[0], 0))
        top = int(round(bbox[1], 0))
        right = int(round(bbox[2], 0))
        bottom = int(round(bbox[3], 0))
        if left < 0:
            left = 0
        if right < 0:
            right = 0
        if top < 0:
            top = 0
        if bottom < 0:
            bottom = 0
        if right > width:
            right = width
        if bottom > height:
            bottom = height
        if left > right:
            left = right
        if top > bottom:
            top = bottom

        if (left == right) or (top == bottom):
            continue

        ret_bbox = [left, top, right, bottom]
        ret_bbox.extend(map(float, bbox[4:]))
        ret_bboxes.append(ret_bbox)

        if show:
            label_text = '{} {:.2f}'.format(class_name, ret_bbox[4])
            cv2.putText(img, label_text, (ret_bbox[0], ret_bbox[1]), cv2.FONT_HERSHEY_COMPLEX, 1.0, color)
            cv2.rectangle(img, (ret_bbox[0], ret_bbox[1]), (ret_bbox[2], ret_bbox[3]), color, thickness=2)
    return ret_bboxes


def margin_crop(bbox, w, h, margins):
    xc = (bbox[0] + bbox[2]) / 2
    yc = (bbox[1] + bbox[3]) / 2
    wbox = (bbox[2] - bbox[0]) / 2
    hbox = (bbox[3] - bbox[1]) / 2

    # Add margin to the patch along each side, based on size from the center of the crop
    xmin = int(round(max(0, xc - wbox * (1 + margins[0]))))
    ymin = int(round(max(0, yc - hbox * (1 + margins[1]))))
    xmax = int(round(min(w, xc + wbox * (1 + margins[2]))))
    ymax = int(round(min(h, yc + hbox * (1 + margins[3]))))

    bbox = (xmin, ymin, xmax, ymax)
    return bbox


def resize_image(bbox_image, scale_size):
    w = bbox_image.shape[1]
    h = bbox_image.shape[0]
    size = max(w, h)
    scale = scale_size / size
    if scale < 1.0:
        w = int(round(w * scale))
        h = int(round(h * scale))
        bbox_image = cv2.resize(bbox_image, (w, h))
    return bbox_image


def crop_proc(file_queue, out_queue, id):
    try:
        score_thr = THRESHOLDS[select_label]

        win_name = 'process{}'.format(id)

        is_break = False
        while True:
            file_info = file_queue.get(block=True)
            if file_info is None:
                break

            file_name, bboxes_infos = file_info
            img_path = os.path.join(base_path, file_name)
            if not os.path.exists(img_path):
                print(img_path, 'no exists')
                continue

            cap = cv2.VideoCapture(img_path)

            image_dir = os.path.splitext(file_name)[0]
            image_path = os.path.join(out_path, image_dir)
            if not os.path.exists(image_path):
                os.makedirs(image_path)

            bbox_infos = []
            many_bbox_infos = []
            frame_id = 0
            while True:
                grabbed, image_bgr = cap.read()

                if not grabbed:
                    break

                bboxes = bboxes_infos.get(frame_id, None)
                if bboxes is not None:
                    bboxes = get_box(image_bgr, bboxes, score_thr=score_thr, show=SHOW)
                    if SHOW:
                        cv2.imshow(win_name, image_bgr)

                    if SAVE_IMAGE:
                        image = image_bgr
                        height = image.shape[0]
                        width = image.shape[1]
                        bbox_count = len(bboxes)
                        many_bboxes = []
                        for i, bbox in enumerate(bboxes):

                            image_name = 'F{:04d}B{:02d}.jpg'.format(frame_id, i)
                            # image_name = 'F{:04d}B{:02d}.png'.format(frame_id, i)

                            mbbox = margin_crop(bbox, width, height, margins)
                            out_image_path = os.path.join(image_path, image_name)
                            bbox_image = image[mbbox[1]:mbbox[3], mbbox[0]:mbbox[2]]
                            bh0 = bbox_image.shape[0]
                            bw0 = bbox_image.shape[1]
                            left = (bbox[0] - mbbox[0])
                            top = (bbox[1] - mbbox[1])
                            right = (bbox[2] - mbbox[0])
                            bottom = (bbox[3] - mbbox[1])
                            bbox_image = resize_image(bbox_image, max_size)
                            bh1 = bbox_image.shape[0]
                            bw1 = bbox_image.shape[1]
                            if bw0 != bw1:
                                left = min(int(round(left * bw1 / bw0)), bw1)
                                right = min(int(round(right * bw1 / bw0)), bw1)
                            if bh0 != bh1:
                                top = min(int(round(top * bh1 / bh0)), bh1)
                                bottom = min(int(round(bottom * bh1 / bh0)), bh1)
                            cv2.imwrite(out_image_path, bbox_image)
                            bbox_info = '{} {},{},{},{},{:.08f} {},{},{},{},{:.08f}\n'.format(image_name, left, top, right, bottom, bbox[4], *bbox)
                            bbox_infos.append(bbox_info)
                            if bbox_count > 1:
                                many_bboxes.append(image_name)
                        if bbox_count > 1:
                            many_bbox_infos.append(many_bboxes)
                else:
                    if SHOW:
                        cv2.imshow(win_name, image_bgr)
                if SHOW:
                    if 0 == SLEEP_TIME:
                        k = cv2.waitKey()
                    else:
                        k = cv2.waitKey(SLEEP_TIME)
                    if (k == 27) or (k == ord('q')) or (k == ord('Q')):  # Esc key to stop
                        is_break = True
                        break
                frame_id += 1

            if SAVE_IMAGE:
                while out_queue.qsize() > 100:
                    time.sleep(0.01)
                out_queue.put(((file_name, bbox_infos, many_bbox_infos), False, id))
            if is_break:
                break
    except Exception as e:
        print('Proc:{}'.format(id), e)
    if SHOW:
        cv2.destroyAllWindows()
    out_queue.put((None, True, id))


def file_proc(file_queue, score_thr):
    with open(bbox_count_file_path, 'r') as file:
        bbox_file = open(bbox_file_path, 'rb')
        file_count = 0
        for line in file.readlines():
            file_count += 1
            lines = line.strip().split(',')
            file_name = lines[0]
            bbox_count = int(lines[1])
            bbox_infos = {}
            for _ in range(bbox_count):
                bbox_info = bbox_file.read(28)
                if len(bbox_info) < 28:
                    for _ in range(PROCESS_COUNT):
                        file_queue.put(None)
                    return
                bbox = struct.unpack('6i1f', bbox_info)
                label = bbox[1]
                if select_label != label:
                    continue
                frame_id = bbox[0]
                score = bbox[6]
                if score < score_thr:
                    continue
                bboxes = bbox_infos.get(frame_id, None)
                if bboxes is None:
                    bbox_infos[frame_id] = [bbox[2:]]
                else:
                    bboxes.append(bbox[2:])

            if len(bbox_infos) > 0:
                while file_queue.qsize() > 100:
                    time.sleep(0.01)
                file_queue.put((file_name, bbox_infos))

            # if file_count >= 10:
            #     break
        bbox_file.close()
    for _ in range(PROCESS_COUNT):
        file_queue.put(None)


def main():
    if not os.path.exists(out_path):
        os.makedirs(out_path)

    file_queue = multiprocessing.Queue()
    out_queue = multiprocessing.Queue()

    score_thr = THRESHOLDS[select_label]

    file_worker = multiprocessing.Process(target=file_proc, args=(file_queue, score_thr))

    file_worker.start()

    select_video_workers = []
    for i in range(PROCESS_COUNT):
        select_video_workers.append(multiprocessing.Process(target=crop_proc, args=(file_queue, out_queue, i,)))

    for i in range(PROCESS_COUNT):
        select_video_workers[i].start()

    out_bbox_count_file_path = os.path.join(out_path, 'bbox_count_list.txt')
    out_bbox_file_path = os.path.join(out_path, 'bbox_list.txt')
    out_bbox_count_file = open(out_bbox_count_file_path, 'w')
    out_bbox_file = open(out_bbox_file_path, 'w')

    out_many_bbox_count_file_path = os.path.join(out_path, 'many_bbox_count_list.txt')
    out_many_bbox_count_file = open(out_many_bbox_count_file_path, 'w')

    file_count = 0
    try:
        finish_worker_count = 0
        while True:
            file_info = out_queue.get(block=True)
            if file_info is None:
                break
            file_bbox_infos, finish, id = file_info
            if finish:
                print('Proc{} finish'.format(id, ))
                finish_worker_count += 1
                if PROCESS_COUNT <= finish_worker_count:
                    break
                continue
            file_name, bbox_infos, many_bbox_infos = file_bbox_infos
            bbox_count = len(bbox_infos)
            if bbox_count > 0:
                for bbox_info in bbox_infos:
                    out_bbox_file.write(bbox_info)
                out_bbox_count_file.write('{},{}\n'.format(file_name, bbox_count))

                if len(many_bbox_infos) > 0:
                    file_base_name = os.path.splitext(file_name)[0]
                    for many_bbox_info in many_bbox_infos:
                        out_many_bbox_count_file.write('{} {}\n'.format(file_base_name, ' '.join(many_bbox_info)))
            file_count += 1
            print('Proc{}, File Count: {}, {}'.format(id, file_count, file_name))
            if file_count % 10 == 0:
                out_bbox_count_file.flush()
                out_bbox_file.flush()
    except Exception as e:
        print(e)

    out_many_bbox_count_file.close()
    out_bbox_count_file.close()
    out_bbox_file.close()
    os.system('chmod a+wr {}'.format(out_many_bbox_count_file_path))
    os.system('chmod a+wr {}'.format(out_bbox_count_file_path))
    os.system('chmod a+wr {}'.format(out_bbox_file_path))

    for i in range(PROCESS_COUNT):
        select_video_workers[i].join()

    file_worker.join()

    os.system('chmod a+wr {} -Rf'.format(out_path))

    print('finish!!!!')


if __name__ == '__main__':
    main()
